#!/usr/bin/env python
# -*- coding: utf-8 -*-

import sys, os, glob
import json
import matplotlib.pyplot as plt
import numpy as np

# 转为mV,CLHEP::volt = 1e-6
millivolt = 1e-9

# 从文件中读取数据
def get(fp):
	fp.seek(0)
	digi_list = [[] for _ in range(50)]
	trace_list = []
	model = {}
	p = 0
	for l in fp.readlines():
		j = json.loads(l)
		if j["class"] == "G4SipmDigi":
			digi_list[p].append(j)
		if j["class"] == "G4SipmVoltageTraceDigi":
			p = p + 1
			trace_list.append(j)
		if j["class"] == "G4SipmModel":
			model = j
	return digi_list,trace_list,model


if __name__ == "__main__":

	# 获取results目录下最新的txt文件
	filename = glob.glob('results/*.txt')[-1]
	if len(sys.argv) > 1:
		filename = sys.argv[1]

	# 从参数中获取起始时间和结束时间
	start = 0
	stop = 1000

	# 打开文件
	fp = open(filename)
	digis, traces, model = get(fp)

	# 解析电压轨迹
	voltages = []
	times = []
	t_signal, t_thermal, t_ct, t_ap = [], [], [], []
	t = []
	labels = []

	tt_voltages = 0
	t_voltages_list = []
	t_voltages_list.append(0)


	# 遍历每个trace中的电压值
	for j, trace in enumerate(traces):
		for q, v in enumerate(trace["voltages"]):
			# 计算时间
			t_voltages = trace["tMin"] + q * trace["timeBinWidth"]
			
			# 如果时间在指定范围内，则将其添加到voltages列表中
			# if t_voltages >= start and t_voltages <= stop:
			voltages.append(v / millivolt)
			times.append((t_voltages + tt_voltages))
		tt_voltages = trace["tMax"] - trace["tMin"] + tt_voltages
		t_voltages_list.append(tt_voltages)
			
	
	print(t_voltages_list)
	# 遍历所有digi中的时间戳
	for i, digi in enumerate(digis):
		for d in digi:
			# d['time'] /= 1000
			if d['type'] == 1:
				t_signal.append(d['time']+t_voltages_list[i])
			elif d['type'] == 2:
				t_thermal.append(d['time']+t_voltages_list[i])
			elif d['type'] == 3:
				t_ct.append(d['time']+t_voltages_list[i])
			elif d['type'] == 4:
				t_ap.append(d['time']+t_voltages_list[i])
	
	
	if len(t_signal) > 0:
		t.append(np.asarray(t_signal))
		labels.append("photon")
	if len(t_thermal) > 0:
		t.append(np.asarray(t_thermal))
		labels.append("thermal")
	if len(t_ct) > 0:
		t.append(np.asarray(t_ct))
		labels.append("crosstalk")
	if len(t_ap) > 0:
		t.append(np.asarray(t_ap))
		labels.append("afterpulse")	

	
	# 计算直方图的bin数和范围
	max_t = max([max(t_i) for t_i in t])
	min_t = min([min(t_i) for t_i in t])
	bins = int((max_t - min_t)/10)
	range = (min_t, max_t)

	fig, ax1 = plt.subplots()
	# 绘制电压波形曲线
	ax1.plot(times, voltages, 'r-')
	ax1.set_ylabel('voltage / mV')
	ax1.set_xlabel(u'time / ns')
	# 创建第二个坐标轴,共享x轴
	ax2 = ax1.twinx()
	colors = "#1f77b4, #ff7f0e, #2ca02c, #d62728".split(", ")[0:len(t)]
	# 绘制直方图，堆叠显示，填充颜色，线宽为0，边缘颜色为白色
	ax2.hist(t, bins=bins, range=range, stacked=True, histtype="stepfilled", color=colors, linewidth=0, ec="white", label=labels)
	ax2.legend(loc="upper right")
	ax2.set_ylabel('entries')
	# 设置图表标题，显示模型名称、尺寸、间距、单元尺寸和单元数量
	plt.text(0.025, 0.975, "%s\n%d x %d mm, %d $\mu$m pitch\n%d cells"  % (model["model"], model["pitch"], model["pitch"], model["cellPitch"] * 1000, model["numberOfCells"]), 
			ha="left", va="top", fontsize="medium", transform=ax1.transAxes)
	# plt.savefig('voltages.svg')
	plt.show()
	
